class Solution:
    def trap(self, height: 'List[int]') -> int:
        if height is None or len(height) <= 2:
            return 0
        res = 0
        # find each height's start & end
        data = {}
        for i, x in enumerate(height):
            if x not in data:
                data[x] = [i, i]
            else:
                data[x] = [data[x][0], i]
        # print(data)

        # get the full state
        sorted_keys = sorted(data, key=lambda x: x, reverse=True)
        for i in range(len(sorted_keys) - 1):
            data[sorted_keys[i+1]][0] = min((data[sorted_keys[i]][0], data[sorted_keys[i+1]][0]))
            data[sorted_keys[i+1]][1] = max((data[sorted_keys[i]][1], data[sorted_keys[i+1]][1]))
        # print(data)

        # resize to height map
        final_height = [0] * len(height)
        # print(final_height)
        for k in reversed(sorted_keys):
            # print('i: ', k)
            for j in range(data[k][0], data[k][1]+1):
                final_height[j] = k
        

        for x, y in zip(height, final_height):
            res += y - x
        return res

if __name__ == "__main__":
    s = Solution()
    print(s.trap([0,1,0,2,1,0]))
    print(s.trap([0,1,0,2,1,0,1,3,2,1,2,1]))
    print(s.trap([2,1]))
    print(s.trap([2,1,3]))